import torch
import torch.nn as nn

"""
nn.LSTM类初始化主要参数解释:
    input_size: 输入张量x中特征维度的大小.
    hidden_size: 隐层张量h中特征维度的大小.
    num_layers: 隐含层的数量.
    bidirectional: 是否选择使用双向LSTM, 如果为True, 则使用; 默认不使用.
"""
# nn.LSTM类实例化对象主要参数解释:
#     input: 输入张量x.
#     h0: 初始化的隐层张量h.
#     c0: 初始化的细胞状态张量c.

# 定义LSTM的参数含义: (input_size, hidden_size, num_layers)
rnn = nn.LSTM(5, 6, 2)
# rnn = nn.LSTM(5, 6, 2, bidirectional=True)
# 定义输入张量的参数含义: (sequence_length, batch_size, input_size)
input001 = torch.randn(2, 3, 5)
# 定义隐藏层初始张量和细胞初始状态张量的参数含义: (num_layers * num_directions, batch_size, hidden_size)
h0 = torch.randn(2, 3, 6)
c0 = torch.randn(2, 3, 6) # 与h0一致

"""使用biLSTM时，层数也要改变"""
# h0 = torch.randn(4, 3, 6)
# c0 = torch.randn(4, 3, 6) # 与h0一致

output001, (h1, c1) = rnn(input001, (h0, c0))
print(output001)
print('-----------------------------')
print(h1)
print('-----------------------------')
print(c1)





